#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#define VIRUS_SIZE 41472
#define INFECTED_SIZE 1000
#define KEY_OFFSET 1234

#define DEBUG(x,args...) // printf(x, ##args)
#define CHECK(fp, in, pos, cmp_val) \
    fseek(fp, pos, SEEK_SET); \
    if(!fread(&in, sizeof(in), 1, fp) || in != (cmp_val)) \
        goto fail

static inline void imul32(uint32_t * dest, uint32_t d, uint32_t s)
{
    int64_t res = (int64_t)(int32_t)d * (int32_t)s;
    *dest = (uint32_t)res;
    return;
}

static inline void mul32(uint32_t * dest, uint32_t d, uint32_t s)
{
    uint64_t res = (uint64_t)d * s;
    *dest = ((uint32_t)(res >> 32));
    return;
}

void bytes_decode(char *data, int count, uint32_t key)
{
    uint32_t res = 0;
    do {
        imul32(&key, key, 0x8088405);
        key++;

        mul32(&res, key, 0xff);
        *data++ ^= res;
    } while(--count);
}

void neshta_fix(const char *name)
{
    uint32_t key;

    FILE *in = fopen(name, "rb");
    if(!in)
    {
        fprintf(stderr, "Failed to open \"%s\": %s\n", name, strerror(errno));
        return;
    }

    fseek(in, 0, SEEK_END);
    size_t len = ftell(in);
    if(len <= 41472) goto fail;

    uint32_t magic;
    CHECK(in, magic, 0x1a89, (uint32_t)0x8088405);
DEBUG("pass1\n");
    uint64_t magic1;
    CHECK(in, magic1, 0xa11b, *((uint64_t*)"Belarus."));
DEBUG("pass2\n");
    printf("Repair[%s]\n", name);

    fseek(in, len - VIRUS_SIZE + KEY_OFFSET, SEEK_SET);
    if(!fread(&key, sizeof(key), 1, in)) goto fail;
DEBUG("pass3\n");
    char *data = (char*)malloc(len - VIRUS_SIZE);
    if(!data)
    {
        fprintf(stderr, "Failed to alloc memory!\n");
        goto fail;
    }
    fseek(in, len - VIRUS_SIZE, SEEK_SET);
DEBUG("%lu\n", len - VIRUS_SIZE + KEY_OFFSET);
    if(!fread(data, VIRUS_SIZE, 1, in)) goto fail;
    bytes_decode(data, INFECTED_SIZE, key);
DEBUG("pass4\n");
    fseek(in, VIRUS_SIZE, SEEK_SET);
    if(!fread(data + VIRUS_SIZE, len - (2 * VIRUS_SIZE), 1, in)) goto fail;
    fclose(in);
DEBUG("pass5\n");
    FILE *out = fopen(name, "wb");
    if(!out)
    {
        fprintf(stderr, "Failed to create file: %s\n", strerror(errno));
        free(data);
        goto fail;
    }
    if(!fwrite(data, len - VIRUS_SIZE, 1, out))
    {
        fprintf(stderr, "Failed to write: %s\n", strerror(errno));
    }
    free(data);
    fflush(out);
    fclose(out);
    return;

fail:
    fclose(in);
    return;
}

int main(int argc, const char ** argv)
{
    if(argc != 2)
    {
        fprintf(stderr, "Usage: %s [infected_exe]\n", argv[0]);
        return 1;
    }

    neshta_fix(argv[1]);
    return 0;
}
